import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Dict, Any, Optional, Union, Tuple
from pathlib import Path

class PandasUtils:
    """Pandas 数据处理工具类，封装常用数据处理、清洗、分析和可视化功能"""
    
    @staticmethod
    def read_data(file_path: str, file_type: Optional[str] = None, **kwargs) -> pd.DataFrame:
        """
        读取不同格式的数据文件
        
        Args:
            file_path: 文件路径
            file_type: 文件类型，可选值：'csv', 'excel', 'json', 'parquet' 等，默认为 None（自动推断）
        
        Returns:
            DataFrame: 读取的数据
        """
        if file_type is None:
            file_ext = Path(file_path).suffix.lower()
            file_type = {
                '.csv': 'csv',
                '.xlsx': 'excel',
                '.xls': 'excel',
                '.json': 'json',
                '.parquet': 'parquet',
                '.h5': 'hdf',
                '.pickle': 'pickle'
            }.get(file_ext, None)
        
        if file_type is None:
            raise ValueError(f"不支持的文件类型: {file_ext}")
        
        readers = {
            'csv': pd.read_csv,
            'excel': pd.read_excel,
            'json': pd.read_json,
            'parquet': pd.read_parquet,
            'hdf': pd.read_hdf,
            'pickle': pd.read_pickle
        }
        
        try:
            return readers[file_type](file_path, **kwargs)
        except Exception as e:
            print(f"读取文件时出错: {e}")
            return pd.DataFrame()
    
    @staticmethod
    def describe_data(df: pd.DataFrame, output_path: Optional[str] = None) -> pd.DataFrame:
        """
        生成数据的描述性统计信息
        
        Args:
            df: 输入数据
            output_path: 输出路径，默认为 None（不输出文件）
        
        Returns:
            DataFrame: 描述性统计信息
        """
        # 基本统计信息
        stats = df.describe(include='all').T
        
        # 添加缺失值信息
        stats['missing_count'] = df.isnull().sum()
        stats['missing_percentage'] = df.isnull().mean() * 100
        
        # 添加数据类型
        stats['dtype'] = df.dtypes
        
        # 添加唯一值数量
        stats['unique_count'] = df.nunique()
        stats['unique_percentage'] = stats['unique_count'] / len(df) * 100
        
        if output_path:
            PandasUtils._save_data(stats, output_path)
        
        return stats
    
    @staticmethod
    def handle_missing_values(df: pd.DataFrame, method: str = 'drop', 
                             columns: Optional[List[str]] = None, 
                             fill_value: Any = None) -> pd.DataFrame:
        """
        处理缺失值
        
        Args:
            df: 输入数据
            method: 处理方法，可选值：'drop', 'fill', 'interpolate'
            columns: 指定列，默认为 None（处理所有列）
            fill_value: 填充值，用于 method='fill'
        
        Returns:
            DataFrame: 处理后的数据
        """
        if columns is None:
            columns = df.columns
        
        if method == 'drop':
            return df.dropna(subset=columns)
        elif method == 'fill':
            if fill_value is None:
                fill_value = {col: df[col].mean() for col in columns if df[col].dtype in [np.number, 'datetime64[ns]']}
                fill_value.update({col: df[col].mode()[0] for col in columns if df[col].dtype not in [np.number, 'datetime64[ns]']})
            return df.fillna(fill_value)
        elif method == 'interpolate':
            return df.interpolate(subset=columns)
        else:
            raise ValueError(f"不支持的方法: {method}")
    
    @staticmethod
    def detect_outliers(df: pd.DataFrame, columns: List[str], method: str = 'iqr', 
                       threshold: float = 1.5) -> pd.DataFrame:
        """
        检测离群值
        
        Args:
            df: 输入数据
            columns: 指定列
            method: 检测方法，可选值：'iqr', 'zscore'
            threshold: 阈值，用于 IQR（1.5 表示中度离群，3 表示极度离群）或 Z-score（通常为 3）
        
        Returns:
            DataFrame: 标记了离群值的布尔型 DataFrame
        """
        outliers = pd.DataFrame(index=df.index)
        
        for col in columns:
            if method == 'iqr':
                q1 = df[col].quantile(0.25)
                q3 = df[col].quantile(0.75)
                iqr = q3 - q1
                lower_bound = q1 - threshold * iqr
                upper_bound = q3 + threshold * iqr
                outliers[col] = (df[col] < lower_bound) | (df[col] > upper_bound)
            elif method == 'zscore':
                z_score = (df[col] - df[col].mean()) / df[col].std()
                outliers[col] = abs(z_score) > threshold
            else:
                raise ValueError(f"不支持的方法: {method}")
        
        return outliers
    
    @staticmethod
    def transform_data(df: pd.DataFrame, transformations: Dict[str, Any]) -> pd.DataFrame:
        """
        对数据进行转换
        
        Args:
            df: 输入数据
            transformations: 转换规则字典，格式为 {列名: 转换函数或参数}
        
        Returns:
            DataFrame: 转换后的数据
        """
        result = df.copy()
        
        for col, transform in transformations.items():
            if col not in result.columns:
                continue
                
            if callable(transform):
                # 应用自定义函数
                result[col] = transform(result[col])
            elif isinstance(transform, dict):
                # 应用内置转换方法
                method = transform.get('method')
                if method == 'log':
                    result[col] = np.log(result[col])
                elif method == 'sqrt':
                    result[col] = np.sqrt(result[col])
                elif method == 'power':
                    power = transform.get('power', 2)
                    result[col] = np.power(result[col], power)
                elif method == 'normalize':
                    result[col] = (result[col] - result[col].mean()) / result[col].std()
                elif method == 'minmax_scale':
                    result[col] = (result[col] - result[col].min()) / (result[col].max() - result[col].min())
            else:
                raise ValueError(f"不支持的转换类型: {transform}")
        
        return result
    
    @staticmethod
    def aggregate_data(df: pd.DataFrame, group_by: Union[str, List[str]], 
                      aggregations: Dict[str, Dict[str, str]]) -> pd.DataFrame:
        """
        数据聚合
        
        Args:
            df: 输入数据
            group_by: 分组列
            aggregations: 聚合规则字典，格式为 {列名: {聚合函数名: 结果列名}}
        
        Returns:
            DataFrame: 聚合结果
        """
        # 转换为 pandas 支持的聚合格式
        agg_spec = {}
        for col, funcs in aggregations.items():
            if col in df.columns:
                agg_spec[col] = list(funcs.keys())
        
        # 执行聚合
        grouped = df.groupby(group_by).agg(agg_spec)
        
        # 重命名列
        if aggregations:
            grouped.columns = ['_'.join(col).strip() for col in grouped.columns.values]
            for col, funcs in aggregations.items():
                for func, new_name in funcs.items():
                    if f"{col}_{func}" in grouped.columns:
                        grouped = grouped.rename(columns={f"{col}_{func}": new_name})
        
        return grouped.reset_index()
    
    @staticmethod
    def merge_data(df1: pd.DataFrame, df2: pd.DataFrame, on: Union[str, List[str]], 
                  how: str = 'inner', suffixes: Tuple[str, str] = ('_x', '_y')) -> pd.DataFrame:
        """
        合并数据
        
        Args:
            df1: 第一个 DataFrame
            df2: 第二个 DataFrame
            on: 连接键
            how: 连接方式，可选值：'inner', 'outer', 'left', 'right'
            suffixes: 后缀，用于区分重复列名
        
        Returns:
            DataFrame: 合并后的结果
        """
        return pd.merge(df1, df2, on=on, how=how, suffixes=suffixes)
    
    @staticmethod
    def pivot_data(df: pd.DataFrame, index: Union[str, List[str]], 
                  columns: str, values: str, aggfunc: str = 'mean') -> pd.DataFrame:
        """
        创建数据透视表
        
        Args:
            df: 输入数据
            index: 行索引
            columns: 列索引
            values: 值
            aggfunc: 聚合函数，默认为 'mean'
        
        Returns:
            DataFrame: 透视表结果
        """
        return df.pivot_table(index=index, columns=columns, values=values, aggfunc=aggfunc)
    
    @staticmethod
    def plot_data(df: pd.DataFrame, x: str, y: Union[str, List[str]], 
                 kind: str = 'line', title: str = '', 
                 figsize: Tuple[int, int] = (10, 6), **kwargs) -> None:
        """
        数据可视化
        
        Args:
            df: 输入数据
            x: x 轴列名
            y: y 轴列名或列名列表
            kind: 图表类型，可选值：'line', 'bar', 'barh', 'scatter', 'hist', 'box', 'pie' 等
            title: 图表标题
            figsize: 图表大小
        """
        plt.figure(figsize=figsize)
        
        if kind == 'pie' and isinstance(y, str):
            # 饼图需要特殊处理
            df.groupby(y)[x].count().plot(kind=kind, autopct='%1.1f%%', **kwargs)
        else:
            df.plot(x=x, y=y, kind=kind, **kwargs)
        
        plt.title(title)
        plt.tight_layout()
        plt.show()
    
    @staticmethod
    def _save_data(df: pd.DataFrame, output_path: str) -> None:
        """
        保存数据到文件
        
        Args:
            df: 要保存的数据
            output_path: 输出路径
        """
        file_ext = Path(output_path).suffix.lower()
        
        if file_ext == '.csv':
            df.to_csv(output_path, index=False)
        elif file_ext in ['.xlsx', '.xls']:
            df.to_excel(output_path, index=False)
        elif file_ext == '.json':
            df.to_json(output_path, orient='records')
        elif file_ext == '.parquet':
            df.to_parquet(output_path, index=False)
        elif file_ext == '.h5':
            df.to_hdf(output_path, key='data', index=False)
        elif file_ext == '.pickle':
            df.to_pickle(output_path)
        else:
            print(f"不支持的输出格式: {file_ext}")

# 使用示例
if __name__ == "__main__":
    # 创建示例数据
    data = {
        'Name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Frank'],
        'Age': [25, 30, 35, 40, np.nan, 22],
        'Salary': [50000, 60000, 75000, 90000, 85000, 45000],
        'Department': ['HR', 'IT', 'IT', 'Finance', 'Finance', 'HR']
    }
    df = pd.DataFrame(data)
    
    # 初始化工具类
    utils = PandasUtils()
    
    # 1. 描述数据
    print("数据描述:")
    print(utils.describe_data(df))
    
    # 2. 处理缺失值
    df_cleaned = utils.handle_missing_values(df, method='fill')
    print("\n处理缺失值后:")
    print(df_cleaned)
    
    # 3. 检测离群值
    outliers = utils.detect_outliers(df_cleaned, columns=['Age', 'Salary'])
    print("\n离群值检测:")
    print(outliers)
    
    # 4. 数据转换
    df_transformed = utils.transform_data(df_cleaned, {
        'Salary': {'method': 'log'},
        'Age': lambda x: x + 1  # 所有年龄加1
    })
    print("\n转换后的数据:")
    print(df_transformed)
    
    # 5. 数据聚合
    aggregated = utils.aggregate_data(df_cleaned, 
                                     group_by='Department', 
                                     aggregations={
                                         'Salary': {'mean': '平均工资', 'max': '最高工资'},
                                         'Age': {'mean': '平均年龄'}
                                     })
    print("\n聚合结果:")
    print(aggregated)
    
    # 6. 数据可视化 (需要图形界面支持)
    # utils.plot_data(df_cleaned, x='Name', y='Salary', kind='bar', title='工资分布')    